﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;

using TracedReality.Logging;
using TracedReality.Primitive;
//using TracedReality.RayTracing;
using TracedReality.SpatialTrees;

namespace TracedReality.RayTracing.PhotonMapping
{
    /// <summary>
    /// The main class for photon mapping.
    /// </summary>
    public class PhotonMap
    {
        private Scene scene;
        private Random random;
        private PhotonTree tree;

        /// <summary>
        /// Create a PhotonMap of the given Scene.
        /// </summary>
        /// <param name="s">The Scene to photon map.</param>
        /// <param name="r">Random to use for sampling.</param>
        public PhotonMap(Scene s, Random r)
        {
            scene = s;
            random = r;
            tree = new PhotonTree(scene.getBoundingBox());
            tracePhotons();
        }

        /// <summary>
        /// Emit and trace Photons from each Light in the Scene.
        /// </summary>
        private void tracePhotons()
        {
            foreach (Light l in scene.lights)
            {
                // This is wrong; each light should shoot a number of photons proportional to its power
                // compared to the other lights
                int np = (int)(Settings.numPhotons / (double)scene.lights.Count);

                for (int a = 0; a < np; a++)
                {
                    tracePhoton(l.generatePhoton(np, random));
                }
            }

            tree.splitCell(Settings.maxPhotonsPerNode, Settings.maxPhotonTreeDepth, random);
            Log.log("Photon tree stats: { Count: " + tree.getCount() + ", Height: " + tree.getHeight() + " }");
        }

        /// <summary>
        /// Recurisvely trace a Photon through the Scene.
        /// Terminates early if the number of bounces is too high.
        /// </summary>
        /// <param name="p">The Photon to trace.</param>
        /// <param name="bounce">How many bounces have occurred.</param>
        private void tracePhoton(Photon p, int bounce = 0)
        {
            if (bounce >= Settings.numPhotonBounces) { return; }

            Ray ray = new Ray(p.position, p.direction);
            Hit hit = new Hit();

            int result = RayTracer.castRay(ray, hit);
            if (result != 0)
            {
                Vector hpos = ray.pointAtParameter(hit.t);
                Vector normal = hit.normal;
                Material m = hit.material;

                if (bounce > 0 && !m.isSpecular())
                {
                    tree.add(new Photon(hpos, p.energy, p.direction));
                }

                double eta1;
                double eta2;
                if (result == 1) { eta1 = Settings.airIndex; eta2 = m.index; }
                else { eta1 = m.index; eta2 = Settings.airIndex; }

                // Find the amount of light reflected and refracted at this angle
                double R;
                if (m.isRefractive())
                {
                    R = Utilities.computeReflectance(p.direction, normal, eta1, eta2);
                }
                else { R = 1; }

                double T = 1 - R;

                // Find the probabilities of diffuse reflection, etc.
                double Pd = m.diffuseProbability(p.energy);
                double Ps = m.reflectiveProbability(p.energy);
                double Pt = m.refractiveProbability(p.energy);
                double s = Pd + (R * Ps) + (T * Pt);
                double r = random.NextDouble();

                Vector rdir = null;
                Vector reng = null;
                Vector npos = null;

                if (r <= Pd)
                {
                    // Diffuse
                    rdir = Vector.RandomDiffuseVector(normal, random);
                    reng = (p.energy * m.diffuse) / Pd;
                    npos = new Vector(hpos);
                }
                else if (r <= (Pd + (R * Ps)))
                {
                    // Reflect
                    rdir = Vector.Reflection(p.direction, normal);
                    reng = (p.energy * m.reflective) / Ps;
                    npos = new Vector(hpos);
                }
                else if (r <= s)
                {
                    // Transmit
                    rdir = Vector.Refraction(p.direction, normal, eta1, eta2);
                    reng = (p.energy * m.refractive) / Pt;
                    npos = hpos - (Settings.EPSILON * normal);
                }
                else
                {
                    // Absorbed
                    return;
                }

                tracePhoton(new Photon(npos, reng, rdir), bounce + 1);
            }
        }

        /// <summary>
        /// Sortable holder for a Photon and its distance to a point.
        /// </summary>
        private class PhotonDistPair : IComparable<PhotonDistPair>
        {
            public Photon photon;
            public double dist;

            public PhotonDistPair(Photon p, double d) { photon = p; dist = d; }

            public int CompareTo(PhotonDistPair other)
            {
                return dist.CompareTo(other.dist);
            }
        }

        /// <summary>
        /// Gather the indirect lighting at a point.
        /// </summary>
        /// <param name="p">The point at which to gather.</param>
        /// <param name="n">The normal of the surface at that point.</param>
        /// <param name="dir">The direction of the incoming Ray.</param>
        /// <returns>The gathered idirect lighting.</returns>
        public Vector gatherIndirect(Vector p, Vector n, Vector dir)
        {
            bool done = false;
            int iteration = 0;
            double radius = Settings.intialGatherRadius;
            List<PhotonDistPair> resulting_photons = new List<PhotonDistPair>();

            // Collect near photons
            do
            {
                resulting_photons.Clear();

                // Get the photons from the tree
                BoundingBox bb = new BoundingBox(p - radius, p + radius);
                List<Photon> near_photons = tree.getInBox(bb);

                // Cull photons that are too far and not in the right direction
                List<PhotonDistPair> culled_photons = new List<PhotonDistPair>();
                foreach (Photon ph in near_photons)
                {
                    if (bb.contains(ph.position))
                    {
                        double d = Vector.Distance(p, ph.position);
                        if (d <= radius && n.dot(-ph.direction) > 0) { culled_photons.Add(new PhotonDistPair(ph, d)); }
                    }
                }

                // Add the n nearest gathered photons to the list
                culled_photons.Sort();
                resulting_photons.AddRange(culled_photons.GetRange(0, Math.Min(100, culled_photons.Count)));

                // If we don't have enough, go to the next iteration
                if (resulting_photons.Count < Settings.numPhotonsToCollect && iteration < Settings.maxGatherIterations)
                {
                    radius += Settings.iterationRadiusIncrease;
                    iteration++;
                }
                else { done = true; }
            } while (!done);

            // Add up all the gathered photons
            Vector gathered_energy = new Vector();
            if (resulting_photons.Count > 0)
            {
                double size = resulting_photons[resulting_photons.Count - 1].dist;
                for (int a = 0; a < resulting_photons.Count; a++)
                {
                    // Cone filter
                    double wgt = (1 - (resulting_photons[a].dist / (Settings.coneFilterK * size)));
                    gathered_energy += resulting_photons[a].photon.energy * wgt;
                }

                // I really don't like that I have to divide by the number collected as well;
                // that was not done in homework 3 but it makes mine look correct
                gathered_energy /= Math.PI * size * size * Settings.numPhotonsToCollect * (1 - (2 / (Settings.coneFilterK * 3)));
            }

            return gathered_energy;
        }
    }
}
